import cv2
import sys
import math
from density_map_process.creat_mark_map import get_warning_map
from model.CrowdCounter import CrowdCounter
from matplotlib import pyplot as plt
from torch.autograd import Variable
from model.misc.utils import *

process_frame_size = (512, 512)
save_frame_size = (400,350)

def create_save_warning_map(et_dmap, image, map_name, threshold, ):

    warning_map = get_warning_map(image, et_dmap, density_threshold= threshold, h = segment, distance=int(distance), sigma_factor = sigma_factor, minimum_area_factor = minimum_area_factor)
    cv2.imwrite(map_name, cv2.resize(warning_map, save_frame_size))

def string_to_float(number):
    try:
        num = float(number)
        return num
    except ValueError:
        print("Cannot convert string to float.")

def closest_power_of_2(num):
    p = int(math.log2(num))
    lower = 2**p
    upper = 2**(p+1)
    if abs(num - lower) < abs(num - upper):
        return lower
    else:
        return upper
mean_std = ([0.452016860247, 0.447249650955, 0.431981861591],[0.23242045939, 0.224925786257, 0.221840232611])
img_transform = standard_transforms.Compose([
        standard_transforms.ToTensor(),
        standard_transforms.Normalize(*mean_std)
    ])

if __name__ == "__main__":


    distance = string_to_float(sys.argv[1])
    sigma_factor = string_to_float(sys.argv[2])
    minimum_area_factor = string_to_float(sys.argv[3])

    segment = string_to_float(sys.argv[4])
    threshold = string_to_float(sys.argv[5])

    if segment > 32:
        segment = 32
    if segment < 2:
        segment = 2
    segment = closest_power_of_2(segment)
    os.chdir("F:\\crowd_density_early_warning_system")
    image_path_root = "images\\warning_images\\image"
    python_stop_flag = "python_stop_flag.txt"
    python_end_flag = "python_end_flag.txt"
    video_path = 'crowd_counting\\metro_trim.mp4'

    model_path = 'crowd_counting\\model_param\\vgg_ep_114_mae_0.9_mse_1.1.pth'  ###更改路径
    device = torch.device("cuda")
    model = CrowdCounter("VGG")
    kpoint = torch.load(model_path, map_location=torch.device('cpu'))
    model.load_state_dict(kpoint)

    # 获取摄像头
    # cap = cv2.VideoCapture(0)
    cap = cv2.VideoCapture(video_path)

    n = 0
    m = 0
    plt.figure(figsize=(5, 5))

    while True:
        m = m + 1
        if m > 50:
            m = 1

        while os.path.exists(python_stop_flag):
            time.sleep(1)
            if os.path.exists(python_end_flag):
                break

        if os.path.exists(python_end_flag):
            break

        # 读取一帧图像
        ret, frame = cap.read()  # fram为读入图片
        img = frame
        if ret == 0:
            break
        cv2.imwrite("images\\crowd_images\\image" + str(int(m)) + ".png", cv2.resize(frame, save_frame_size))

        if img.shape[0] > 0 and img.shape[1] > 0:
            n = n + 1
            if n / 5 > 50:
                n = 1
            if n % 5 == 0:

                img = cv2.resize(img,process_frame_size, interpolation=cv2.INTER_CUBIC)
                img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
                img = img_transform(img)

                with torch.no_grad():
                    img = Variable(img[None, :, :, :])
                    pred_map = model.test_forward(img)
                pred_map = pred_map.cpu().data.numpy()[0, 0, :, :]
                pred = str(int(np.sum(pred_map) / 100.0))

                print("当前预测人数：" + pred)
                create_save_warning_map(pred_map/100,
                                        cv2.resize(frame, process_frame_size),
                                        image_path_root + str(int(n / 5)) + ".png",
                                        threshold)

cap.release()
cv2.destroyAllWindows()